"""Encrypt SSH and certificates private keys

Revision ID: 2ab85192c49f
Revises: 83dae64197f3
Create Date: 2020-04-29 18:10:31.488781+00:00

"""
from alembic import op
import sqlalchemy as sa

from middlewared.plugins.pwenc import encrypt


# revision identifiers, used by Alembic.
revision = '2ab85192c49f'
down_revision = '83dae64197f3'
branch_labels = None
depends_on = None


def upgrade():
    conn = op.get_bind()
    if not (
        conn.execute("SELECT * FROM system_keyvalue WHERE key = 'has_0039_auto_20200429_0631' AND value = 'true'").
                fetchall()
    ):
        for table, fields in [
            ("services_snmp", ["snmp_v3_password", "snmp_v3_privpassphrase"]),
            ("services_ssh", ["ssh_privatekey", "ssh_host_dsa_key", "ssh_host_ecdsa_key", "ssh_host_ed25519_key",
                              "ssh_host_key", "ssh_host_rsa_key"]),
            ("services_s3", ["s3_secret_key"]),
            ("system_certificate", ["cert_privatekey"]),
            ("system_certificateauthority", ["cert_privatekey"]),
        ]:
            for row in conn.execute(f"SELECT * FROM {table}").fetchall():
                set_ = []
                params = []
                for k in fields:
                    if row[k] is not None:
                        set_.append(f"{k} = ?")
                        params.append(encrypt(row[k]))

                if set_:
                    conn.execute(f"UPDATE {table} SET {', '.join(set_)} WHERE id = {row['id']}", params)

    conn.execute("DELETE FROM system_keyvalue WHERE key = 'has_0039_auto_20200429_0631'")


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    pass
    # ### end Alembic commands ###
